
try:
    from detect_evaluate import np, cvio, os, calculate_overlaps
except Exception as e:
    print(e)
    from .detect_evaluate import np, cvio, os, calculate_overlaps

def load_dataset(src):
    gtlist = cvio.load_ext_list(src, ext_type='.json')
    dataset = []
    for gt in gtlist:
        gta = cvio.load_ann(gt)
        dataset.append([gt, gta])
    return dataset


def filter_by_spec_contour(ann, spec='a', iou_thr=0.99, mode='bbox', iou_mode='minimum'):
    spec_shapes = []
    other_shapes = []
    for shape in ann['shapes']:
        label = shape['label']
        if label != spec:
            other_shapes.append(shape)
        else:
            spec_shapes.append(shape)
    
    ovlps = calculate_overlaps(spec_shapes, other_shapes, mode=mode, iou_mode=iou_mode)
    if not len(ovlps):
        return ann
    keep_shapes = []
    for i, iou in enumerate(ovlps):
        # print(i, iou)
        ids = np.where(iou >iou_thr)[0]
        keep_shapes.extend([other_shapes[j] for j in ids])

    ann['shapes'] = keep_shapes

    return ann

def run_keep_main_ann(src, dst, spec='a', iou_thr=0.99, mode='mask', iou_mode='minimum'):
    anns = load_dataset(src)
    n = len(anns)
    for i, (path, ann) in enumerate(anns, 1):
        print('[%d/%d]' % (i, n), path)
        ann = filter_by_spec_contour(
            ann, spec=spec, iou_thr=iou_thr, mode=mode,iou_mode=iou_mode)
        cvio.write_ann(ann, os.path.join(dst, os.path.basename(path)))

if __name__ == "__main__":
    src = r'J:\fmg_test\fmg_test\merge_a_p'
    dst = r'J:\fmg_test\fmg_test\after'
    run_keep_main_ann(src, dst, spec='a', iou_thr=0.8)
